#!/usr/bin/env python3

import argparse
import os
import glob
import numpy as np
from astropy.io import fits


def raw2fits(path, suffix='raw', row=5000, col=6000,
             overwrite=False):
    # check if input is directory or file name
    if os.path.exists(path):
        if os.path.isdir(path):
            flist = glob.glob('{}{}*.{}'.format(path, os.sep, suffix))
        else:
            flist = [path]
    else:
        raise Exception('invalid path: {}'.format(path))

    for f in flist:
        # check if already converted
        oname = f.replace('.raw', '.fits')
        if os.path.exists(oname) and (not overwrite):
            continue
        # read data
        data = np.fromfile(f, dtype='<u4')
        data = data.reshape(row, col)
        hdu = fits.PrimaryHDU(data)
        # extract information from file name
        hdu.writeto(oname, overwrite=True)
        print('File converted: {}'.format(oname))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('path', type=str)
    parser.add_argument('--suffix', default='raw', type=str)
    parser.add_argument('--row', default=5000, type=int)
    parser.add_argument('--col', default=6000, type=int)
    parser.add_argument('--overwrite', action='store_true')
    args = parser.parse_args()
    raw2fits(args.path, suffix=args.suffix, row=args.row, col=args.col,
             overwrite=args.overwrite)
